import email
import re
import base64
from email import policy
from email.parser import BytesParser


def process_eml_with_images(eml_content):
    """
    解析EML内容并替换其中的CID图片引用为Base64编码的图片

    :param eml_content: EML文件的二进制内容
    :return: 处理后的HTML内容
    """
    # 解析EML文件
    msg = BytesParser(policy=policy.default).parsebytes(eml_content)

    # 获取HTML内容
    html = get_html_body(msg)
    if not html:
        text = get_text_body(msg)
        html = text.replace('\n', '<br>') if text else ''

    # 获取所有内联图片的CID映射
    cid_map = get_cid_map(msg)

    # 替换HTML中的CID引用为Base64编码的图片
    html = replace_cid_with_base64(html, cid_map)

    return html


def get_html_body(msg):
    """从邮件中提取HTML正文"""
    if msg.is_multipart():
        for part in msg.walk():
            content_type = part.get_content_type()
            if content_type == 'text/html':
                payload = part.get_payload(decode=True)
                charset = part.get_content_charset() or 'utf-8'
                return payload.decode(charset, errors='replace')
    else:
        if msg.get_content_type() == 'text/html':
            payload = msg.get_payload(decode=True)
            charset = msg.get_content_charset() or 'utf-8'
            return payload.decode(charset, errors='replace')
    return None


def get_text_body(msg):
    """从邮件中提取纯文本正文"""
    if msg.is_multipart():
        for part in msg.walk():
            content_type = part.get_content_type()
            if content_type == 'text/plain':
                payload = part.get_payload(decode=True)
                charset = part.get_content_charset() or 'utf-8'
                return payload.decode(charset, errors='replace')
    else:
        if msg.get_content_type() == 'text/plain':
            payload = msg.get_payload(decode=True)
            charset = msg.get_content_charset() or 'utf-8'
            return payload.decode(charset, errors='replace')
    return None


def get_cid_map(msg):
    """创建CID到图片数据的映射"""
    cid_map = {}

    if msg.is_multipart():
        for part in msg.walk():
            content_disposition = part.get('Content-Disposition', '')
            if 'inline' in content_disposition or part.get_content_type().startswith('image/'):
                cid = part.get('Content-ID', '')
                if cid:
                    cid = cid.strip('<>')
                    image_data = part.get_payload(decode=True)
                    content_type = part.get_content_type()
                    cid_map[cid] = {
                        'data': image_data,
                        'type': content_type
                    }
    return cid_map


def replace_cid_with_base64(html, cid_map):
    """替换HTML中的CID引用为Base64编码的图片"""

    def replace_match(match):
        cid = match.group(1)
        if cid in cid_map:
            image_data = cid_map[cid]['data']
            content_type = cid_map[cid]['type']
            base64_data = base64.b64encode(image_data).decode('ascii')
            return f'<img src="data:{content_type};base64,{base64_data}" />'
        return match.group(0)

    # 替换cid引用
    pattern = re.compile(r'<img[^>]+src=["\']cid:([^"\']+)["\'][^>]*>', re.IGNORECASE)
    html = pattern.sub(replace_match, html)

    return html


# 使用示例
if __name__ == '__main__':
    # 读取EML文件
    with open('example2.eml', 'rb') as f:
        eml_content = f.read()

    # 处理EML内容
    processed_html = process_eml_with_images(eml_content)

    # 输出结果
    print(processed_html)